#pragma once

#include <brpc/server.h>
#include <brpc/channel.h>
#include <unistd.h>
#include <butil/logging.h>
#include "dmatest.pb.h"
#include "doca_dma_unit.h"
#include "kv.pb.h"

DEFINE_string(protocol, "baidu_std", "Protocol type. Defined in src/brpc/options.proto");
DEFINE_string(connection_type, "", "Connection type. Available values: single, pooled, short");
DEFINE_int32(timeout_ms, 10000, "RPC timeout in milliseconds");

#if defined(__aarch64__)
DEFINE_string(addr_and_port, "10.10.226.50:31010", "server brpc port");
#else
DEFINE_string(addr_and_port, "10.10.226.51:31010", "server brpc port");
#endif

// DEFINE_string(addr_and_port, "127.0.0.1:31111", "server brpc port");
DEFINE_string(load_balancer, "", "The algorithm for load balancing");

namespace status_type{
    enum mmap_type{
        COME,
        TO
    };
};

/*专门做DMA， TODO: 搞一个base的brpc通信unit，初始化把这些protocol这种都赋值好*/ 
class BrpcUnit{
public:
    static BrpcUnit & GetInstance() {
        static BrpcUnit brpcunit;
        return brpcunit;
    }

    void InitDmaUnit(const std::vector<size_t> & offsets){
        dmaunit = 
            new DmaUnit(
                0,
                0,
                1, // 看是否创建回调函数, 1代表是
                (int)offsets.size() - 1,
                offsets
            );
    }

    void InitDmaUnit(const std::vector<size_t> & offsets, void * ptr){
        dmaunit = 
            new DmaUnit(
                0,
                0,
                1, // 看是否创建回调函数, 1代表是
                (int)offsets.size() - 1,
                offsets,
                ptr
            );
    }

    void * get_ptr() {
        return dmaunit->get_ptr();
    }

    void task_dispatch_mem(const std::vector<int>& task_mem_list) {
        if ( (int)task_mem_list.size() == 0) return;
        dmaunit->StartDmaCopy(task_mem_list);
    }

    void task_dispatch_op(const std::vector<int>& task_op_list, int layer_id = 0) {
        rap::DocaDmaService_Stub rpc_chan(&this->chan);
        brpc::Controller cntl;

        rap::TaskSend request;
        rap::TaskFinish response;

        for (int node_id : task_op_list) {
            request.add_operation_node_list(node_id);  // 逐个添加元素
        }
        request.set_layer_id(layer_id);

        rpc_chan.TaskDispatch(&cntl, &request, &response, NULL);
        if (cntl.Failed()) {
            LOG(ERROR) << "task brpc transfer failed";
        }
    }

    void export_info (const rap::DmaUnitInit * req, rap::DmaUnitResponse * resp) {
        // 从client发过来的返回的export信息。
        std::string export_desc_from(req->export_descs().data(), req->export_descs().size());
        dmaunit->ExportFromDesc(export_desc_from);

        //发送当前的export信息到调度器
        std::string export_desc_to = dmaunit->GetExportDesc(); 
        resp->set_export_descs(export_desc_to.data(), export_desc_to.size() );
    }

    void InitBrpc() {
        options.protocol = FLAGS_protocol;
        options.connection_type = FLAGS_connection_type;
        options.timeout_ms = FLAGS_timeout_ms /*milliseconds*/;

        if (chan.Init(FLAGS_addr_and_port.data(), FLAGS_load_balancer.c_str(), &options) != 0) {
            LOG(ERROR) << "Fail to initialize channel";
        }
    }

    void SendDmaInfo() {      
        LOG(INFO) << "start send doca_init request";
        std::string export_descs = dmaunit->GetExportDesc();
        rap::DocaDmaService_Stub rpc_chan(&this->chan);
        brpc::Controller cntl;

        rap::DmaUnitInit init_request;
        rap::DmaUnitResponse init_response;

        init_request.set_export_descs(export_descs.data(), export_descs.size());

        rpc_chan.InitDmaChannel(&cntl, &init_request, &init_response, NULL);
        if (cntl.Failed()) {
            LOG(ERROR) << "brpc init connect failed";
        }
        // 判断是否成功
        std::string export_desc_from(init_response.export_descs().data(), init_response.export_descs().size());
        dmaunit->ExportFromDesc(export_desc_from);
    }

    /* 用于单纯的dma通信创建，普通的根据发送的各类信息：包括task数量，每个task的偏移等信息创建dmaunit，根据export descs创建memcpy-task*/
    void InitDma (const rap::DmaUnitInit * req, rap::DmaUnitResponse * resp) {
        std::vector<size_t> offsets;
        for (int i = 0; i < req->offsets_size(); ++i) {
            offsets.push_back(req->offsets(i));
        }
        if (dmaunit != nullptr) {
            delete dmaunit;
        }
        dmaunit = 
            new DmaUnit(
                req->device_type(),
                req->device_id(),
                ( (req->mmap_type() ) ^ 1), // 看是否创建回调函数
                req->task_num(),
                offsets
            );
        
        if ( (req->mmap_type()) == 0) {
            std::string export_desc(req->export_descs().data(), req->export_descs().size());
            dmaunit->ExportFromDesc(export_desc);
        }
        else {
            std::string export_desc = dmaunit->GetExportDesc();
            resp->set_export_descs(export_desc.data(), export_desc.size() );
        }
    }
    /* 普通的dma通信，发送详细的信息到对面*/
    void SendDmaInfo(int device_type, int device_id, int mmap_type, int task_num, const std::vector<size_t>& offsets) {
        std::string export_descs = dmaunit->GetExportDesc();
        rap::DocaDmaService_Stub rpc_chan(&this->chan);
        brpc::Controller cntl;

        rap::DmaUnitInit init_request;
        rap::DmaUnitResponse init_response;

        // 设置新的参数
        init_request.set_device_type(device_type);
        init_request.set_device_id(device_id);
        init_request.set_mmap_type(mmap_type);
        init_request.set_task_num(task_num);
        for (const auto& offset : offsets) {
            init_request.add_offsets(offset);
        }
        init_request.set_export_descs(export_descs.data(), export_descs.size());

        rpc_chan.InitDmaChannel(&cntl, &init_request, &init_response, NULL);
        if (cntl.Failed()) {
            LOG(ERROR) << "brpc init connect failed";
        }

        std::string export_desc_from(init_response.export_descs().data(), init_response.export_descs().size());
        dmaunit->ExportFromDesc(export_desc_from);
    }

public:
    brpc::Channel chan;
    brpc::ChannelOptions options;

public:
    DmaUnit * dmaunit = nullptr;
};

class BrpcKvstore {
public:

    class PromiseWrapper {
    public:
        PromiseWrapper() : promise_(make_unique<std::promise<int>>()) {}
            std::future<int> GetFuture() {
            return promise_->get_future();
        }
        void SetValue(int value) {
            promise_->set_value(value);
        }
        void Reset() {
            promise_ = make_unique<std::promise<int>>();
        }

    private:
        std::unique_ptr<std::promise<int>> promise_;
    };

    static BrpcKvstore& getInstance() {
        static BrpcKvstore instance;
        return instance;
    }

    void InitChannel() {
        options_.protocol = "baidu_std";
        options_.timeout_ms = 10000 /*milliseconds*/;
        // if (chan_.Init("10.50.82.188:31011", "", &options_) != 0) {
        if (chan_.Init("127.0.0.1:31011", "", &options_) != 0) {
            LOG(ERROR) << "Fail to initialize channel";
        }
        rpc_chan_ = new ks::reco::DocaDpuService_Stub(dynamic_cast<google::protobuf::RpcChannel*>(&chan_));
    }

    void Init_kv_store(int embedding_dim, std::string DataType, std::vector<std::string> table_names, std::vector<int> num_embeddings) {
        brpc::Controller cntl;
        cntl.set_timeout_ms(1000);

        ks::reco::Initkv request;
        ks::reco::Initresponse response;

        request.set_datatype(DataType);
        request.set_embedding_dim(embedding_dim);

        printf ("you already bind c++ to py\n");
        int num_tables = (int)table_names.size();
        printf ("your table_length = %d\n", num_tables);
        if (num_tables == 1) printf ("table_name : %s\n", table_names[0].data());

        for (int i = 0; i < num_tables; ++i) {
            request.add_table_names(table_names[i]);
            request.add_num_embeddings(num_embeddings[i]);
        }

        if (rpc_chan_ == nullptr) printf ("null rpc_chan_\n");

        rpc_chan_->Initkvstore(dynamic_cast<google::protobuf::RpcController*>(&cntl), &request, &response, NULL);
        if (cntl.Failed()) {
            printf ("connect failed\n");
        }
        if (response.status() != 1) {
            printf ("init_kv_store failed\n");
        }
        else printf ("init_kv_store success\n");
    }

    void handle_embedding_response(
        brpc::Controller* cntl,
        ks::reco::DpuFetchResponse* fetch_response,
        void* output_ptr, 
        int element_size,   // size of each element in bytes
        int num_keys,       // output的行数
        int embedding_size, // output的列数
        int Start_key,
        int End_key,
        PromiseWrapper* promise_ptr
    ) {
        (void) End_key;
        std::unique_ptr<brpc::Controller> cntl_guard(cntl);
        std::unique_ptr<ks::reco::DpuFetchResponse> response_guard(fetch_response);

        int received_key_size = fetch_response->embeddings_size();

        // 检查output的形状是否正确
        if (num_keys < Start_key + received_key_size) {
            throw std::runtime_error("Output numpy array shape mismatch");
        }

        // 将数据写入output数组
        for (int j = 0; j < received_key_size; ++j) {
            auto& embedding_data = fetch_response->embeddings(j);
            std::memcpy(
                static_cast<uint8_t*>(output_ptr) + (Start_key + j) * embedding_size * element_size,
                embedding_data.data(),
                embedding_data.size()
            );
        }

        // 设置Promise的值
        promise_ptr->SetValue(1);
    }

    template <typename Tinput, typename Toutput>
    void fetch_embedding_async(
        const std::string& table_name,   // embedding table name
        int split_num,                   // number of async requests
        Tinput * input_data,          // pointer to input data
        Toutput* output_data,               // pointer to output data
        int total_key_num,               // total number of keys
        int embedding_size              // size of each embedding
    ) {
        int split_times = (total_key_num - 1) / split_num + 1;
        std::vector<PromiseWrapper> promise_wrappers_(split_times);
        int split_cnt = 0;

        for (int i = 0; i < total_key_num; i += split_num, split_cnt++) {
            int End = i + split_num > total_key_num ? total_key_num : i + split_num;

            ks::reco::DpuFetchRequest fetch_request;
            ks::reco::DpuFetchResponse* fetch_response = new ks::reco::DpuFetchResponse();
            brpc::Controller* cntl = new brpc::Controller();

            cntl->set_timeout_ms(1000);
            fetch_request.set_table_name(table_name);
            fetch_request.set_option_type("get");

            // Add keys to the request
            for (int j = i; j < End; ++j) {
                Tinput* input_ptr = reinterpret_cast<Tinput*>(
                    reinterpret_cast<uint8_t*>(input_data) + j * sizeof(Tinput)
                );
                fetch_request.add_key(*input_ptr);
            }

            promise_wrappers_[split_cnt].Reset();

            google::protobuf::Closure* done = brpc::NewCallback(
                this,
                &BrpcKvstore::handle_embedding_response,
                cntl,
                fetch_response,
                static_cast<void*>(output_data),
                (int)sizeof(Toutput),
                total_key_num,
                embedding_size,
                i, End,
                &promise_wrappers_[split_cnt]
            );

            rpc_chan_->DpueFetchByTable(
                dynamic_cast<google::protobuf::RpcController*>(cntl),
                &fetch_request,
                fetch_response,
                done
            );
        }

        for (int i = 0; i < split_times; ++i) {
            auto future = promise_wrappers_[i].GetFuture();
            future.get();
        }
    }

    
public:
    bool enable_fp16 = true;
    brpc::Channel chan_;
    brpc::ChannelOptions options_;
    ks::reco::DocaDpuService_Stub *rpc_chan_;
    std::vector<PromiseWrapper> promise_wrappers_; // for batch_wait
    std::vector<std::vector<float>> embeddings; // fetch all embeddings
};